package com.dmz.common.utils.secure;

import org.bouncycastle.jce.provider.BouncyCastleProvider;

import javax.crypto.Cipher;
import javax.crypto.KeyGenerator;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.SecretKeySpec;
import java.security.NoSuchAlgorithmException;
import java.security.NoSuchProviderException;
import java.security.SecureRandom;
import java.security.Security;
import java.util.Base64;

/**
 * @author lee
 */
public class SM4Util {
    static {
        Security.addProvider(new BouncyCastleProvider());
    }

    private static final String ALGORITHM_NAME = "SM4";
    private static final String ALGORITHM_NAME_ECB_PADDING = "SM4/ECB/PKCS5Padding";
    private static final String ALGORITHM_NAME_CBC_PADDING = "SM4/CBC/PKCS5Padding";
    private static final int DEFAULT_KEY_SIZE = 128;

    /**
     * 生成SM4密钥
     */
    public static byte[] generateKey() throws NoSuchAlgorithmException, NoSuchProviderException {
        KeyGenerator kg = KeyGenerator.getInstance(ALGORITHM_NAME, BouncyCastleProvider.PROVIDER_NAME);
        kg.init(DEFAULT_KEY_SIZE, new SecureRandom());
        return kg.generateKey().getEncoded();
    }

    /**
     * ECB模式加密
     */
    public static byte[] encryptEcb(byte[] key, byte[] data) throws Exception {
        Cipher cipher = Cipher.getInstance(ALGORITHM_NAME_ECB_PADDING, BouncyCastleProvider.PROVIDER_NAME);
        SecretKeySpec sm4Key = new SecretKeySpec(key, ALGORITHM_NAME);
        cipher.init(Cipher.ENCRYPT_MODE, sm4Key);
        return cipher.doFinal(data);
    }

    /**
     * ECB模式解密
     */
    public static byte[] decryptEcb(byte[] key, byte[] cipherText) throws Exception {
        Cipher cipher = Cipher.getInstance(ALGORITHM_NAME_ECB_PADDING, BouncyCastleProvider.PROVIDER_NAME);
        SecretKeySpec sm4Key = new SecretKeySpec(key, ALGORITHM_NAME);
        cipher.init(Cipher.DECRYPT_MODE, sm4Key);
        return cipher.doFinal(cipherText);
    }

    /**
     * CBC模式加密
     */
    public static byte[] encryptCbc(byte[] key, byte[] iv, byte[] data) throws Exception {
        Cipher cipher = Cipher.getInstance(ALGORITHM_NAME_CBC_PADDING, BouncyCastleProvider.PROVIDER_NAME);
        SecretKeySpec sm4Key = new SecretKeySpec(key, ALGORITHM_NAME);
        IvParameterSpec ivParameterSpec = new IvParameterSpec(iv);
        cipher.init(Cipher.ENCRYPT_MODE, sm4Key, ivParameterSpec);
        return cipher.doFinal(data);
    }

    /**
     * CBC模式解密
     */
    public static byte[] decryptCbc(byte[] key, byte[] iv, byte[] cipherText) throws Exception {
        Cipher cipher = Cipher.getInstance(ALGORITHM_NAME_CBC_PADDING, BouncyCastleProvider.PROVIDER_NAME);
        SecretKeySpec sm4Key = new SecretKeySpec(key, ALGORITHM_NAME);
        IvParameterSpec ivParameterSpec = new IvParameterSpec(iv);
        cipher.init(Cipher.DECRYPT_MODE, sm4Key, ivParameterSpec);
        return cipher.doFinal(cipherText);
    }

    /**
     * 生成Base64编码的密钥
     */
    public static String generateKeyBase64() throws NoSuchAlgorithmException, NoSuchProviderException {
        return Base64.getEncoder().encodeToString(generateKey());
    }

    /**
     * 从Base64编码的密钥字符串加载密钥
     */
    public static byte[] loadKeyBase64(String keyBase64) {
        return Base64.getDecoder().decode(keyBase64);
    }
}